import os
import hashlib
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from PIL import Image, ImageTk
import imagehash
from collections import defaultdict
import threading

class DuplicateImageFinder:
    def __init__(self, root):
        self.root = root
        self.root.title("重复图片查找器")
        self.root.geometry("1400x900")
        
        # 存储相似图片组
        self.similar_groups = []
        self.current_group_index = 0
        
        self.setup_ui()
        
    def setup_ui(self):
        # 主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 配置网格权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)
        main_frame.rowconfigure(2, weight=1)
        
        # 文件夹选择
        ttk.Label(main_frame, text="选择文件夹:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.folder_path = tk.StringVar()
        ttk.Entry(main_frame, textvariable=self.folder_path, width=50).grid(row=0, column=1, sticky=(tk.W, tk.E), padx=5)
        ttk.Button(main_frame, text="浏览", command=self.browse_folder).grid(row=0, column=2, padx=5)
        
        # 控制按钮
        control_frame = ttk.Frame(main_frame)
        control_frame.grid(row=1, column=0, columnspan=3, pady=10)
        
        ttk.Button(control_frame, text="开始扫描", command=self.start_scan).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_frame, text="上一组", command=self.prev_group).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_frame, text="下一组", command=self.next_group).pack(side=tk.LEFT, padx=5)
        ttk.Button(control_frame, text="删除当前组小文件", command=self.delete_smaller_files).pack(side=tk.LEFT, padx=5)
        
        # 进度条
        self.progress = ttk.Progressbar(main_frame, mode="indeterminate")
        self.progress.grid(row=1, column=3, sticky=(tk.W, tk.E), padx=5)
        
        # 结果显示区域
        result_frame = ttk.LabelFrame(main_frame, text="相似图片组", padding="5")
        result_frame.grid(row=2, column=0, columnspan=4, sticky=(tk.W, tk.E, tk.N, tk.S), pady=5)
        result_frame.columnconfigure(0, weight=1)
        result_frame.rowconfigure(0, weight=1)
        
        # 创建滚动区域
        canvas = tk.Canvas(result_frame)
        scrollbar = ttk.Scrollbar(result_frame, orient="vertical", command=canvas.yview)
        self.scrollable_frame = ttk.Frame(canvas)
        
        self.scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=self.scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        canvas.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))
        
        # 状态标签
        self.status_label = ttk.Label(main_frame, text="请选择文件夹开始扫描")
        self.status_label.grid(row=3, column=0, columnspan=4, pady=5)
        
    def browse_folder(self):
        folder = filedialog.askdirectory()
        if folder:
            self.folder_path.set(folder)
            
    def start_scan(self):
        folder = self.folder_path.get()
        if not folder:
            messagebox.showwarning("警告", "请先选择文件夹")
            return
            
        # 在新线程中执行扫描
        self.progress.start()
        self.status_label.config(text="正在扫描图片...")
        
        thread = threading.Thread(target=self.scan_images, args=(folder,))
        thread.daemon = True
        thread.start()
        
    def scan_images(self, folder):
        try:
            # 支持的图片格式
            image_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff", ".webp"}
            
            # 收集所有图片文件
            image_files = []
            for root, dirs, files in os.walk(folder):
                for file in files:
                    if os.path.splitext(file.lower())[1] in image_extensions:
                        image_files.append(os.path.join(root, file))
            
            if not image_files:
                self.root.after(0, lambda: self.show_message("未找到图片文件"))
                return
                
            # 计算图片哈希值
            hash_dict = defaultdict(list)
            total_files = len(image_files)
            
            for i, image_path in enumerate(image_files):
                try:
                    # 计算感知哈希
                    with Image.open(image_path) as img:
                        # 转换为RGB模式
                        if img.mode != "RGB":
                            img = img.convert("RGB")
                        phash = imagehash.phash(img)
                        hash_dict[phash].append(image_path)
                        
                    # 更新进度
                    progress = (i + 1) / total_files * 100
                    self.root.after(0, lambda p=progress: self.status_label.config(
                        text=f"正在处理图片... {p:.1f}%"
                    ))
                    
                except Exception as e:
                    print(f"处理图片失败 {image_path}: {e}")
                    continue
            
            # 找出相似图片组
            self.similar_groups = []
            for phash, files in hash_dict.items():
                if len(files) > 1:
                    # 按文件大小排序
                    files_with_size = []
                    for file_path in files:
                        try:
                            size = os.path.getsize(file_path)
                            files_with_size.append((file_path, size))
                        except:
                            continue
                    
                    files_with_size.sort(key=lambda x: x[1], reverse=True)
                    self.similar_groups.append(files_with_size)
            
            # 按组大小排序
            self.similar_groups.sort(key=lambda x: len(x), reverse=True)
            
            self.root.after(0, self.scan_complete)
            
        except Exception as e:
            self.root.after(0, lambda: self.show_error(f"扫描失败: {str(e)}"))
        finally:
            self.root.after(0, self.progress.stop)
            
    def scan_complete(self):
        self.progress.stop()
        if not self.similar_groups:
            self.status_label.config(text="未找到重复图片")
            self.clear_results()
        else:
            self.status_label.config(text=f"找到 {len(self.similar_groups)} 组相似图片")
            self.current_group_index = 0
            self.display_current_group()
            
    def display_current_group(self):
        # 清除之前的结果
        for widget in self.scrollable_frame.winfo_children():
            widget.destroy()
            
        if not self.similar_groups:
            return
            
        group = self.similar_groups[self.current_group_index]
        
        # 显示组信息
        group_info = ttk.Label(
            self.scrollable_frame, 
            text=f"第 {self.current_group_index + 1} 组 (共 {len(self.similar_groups)} 组) - {len(group)} 张相似图片",
            font=("Arial", 12, "bold")
        )
        group_info.grid(row=0, column=0, columnspan=4, pady=10)
        
        # 显示图片
        for i, (file_path, file_size) in enumerate(group):
            self.display_image(file_path, file_size, i + 1, i == 0)  # 第一个是最大的文件
            
    def display_image(self, file_path, file_size, index, is_largest=False):
        try:
            # 创建图片框架
            frame_title = f"图片 {index}"
            if is_largest:
                frame_title += " (最大文件 - 将保留)"
            
            img_frame = ttk.LabelFrame(self.scrollable_frame, text=frame_title)
            img_frame.grid(row=index, column=0, sticky=(tk.W, tk.E), padx=5, pady=5)
            
            # 创建图片和信息的水平布局
            content_frame = ttk.Frame(img_frame)
            content_frame.grid(row=0, column=0, sticky=(tk.W, tk.E), padx=5, pady=5)
            
            # 加载并调整图片大小
            with Image.open(file_path) as img:
                # 计算缩略图大小
                max_size = 200
                img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
                
                # 转换为PhotoImage
                photo = ImageTk.PhotoImage(img)
                
                # 显示图片
                img_label = ttk.Label(content_frame, image=photo)
                img_label.image = photo  # 保持引用
                img_label.grid(row=0, column=0, padx=5, pady=5)
                
                # 显示文件信息
                info_text = f"文件: {os.path.basename(file_path)}\n"
                info_text += f"路径: {file_path}\n"
                info_text += f"大小: {self.format_file_size(file_size)}"
                
                info_label = ttk.Label(content_frame, text=info_text, justify=tk.LEFT)
                info_label.grid(row=0, column=1, padx=5, pady=5, sticky=tk.W)
                
                # 添加删除按钮（除了最大的文件）
                if not is_largest:
                    delete_btn = ttk.Button(
                        content_frame, 
                        text="删除此文件", 
                        command=lambda: self.delete_single_file(file_path, index),
                        style="Accent.TButton"
                    )
                    delete_btn.grid(row=0, column=2, padx=5, pady=5)
                else:
                    # 为最大文件显示保留标签
                    keep_label = ttk.Label(content_frame, text="保留", foreground="green", font=("Arial", 10, "bold"))
                    keep_label.grid(row=0, column=2, padx=5, pady=5)
                
        except Exception as e:
            error_label = ttk.Label(
                self.scrollable_frame, 
                text=f"无法加载图片 {file_path}: {str(e)}"
            )
            error_label.grid(row=index, column=0, sticky=(tk.W, tk.E), padx=5, pady=5)
            
    def delete_single_file(self, file_path, index):
        """删除单个文件"""
        try:
            if os.path.exists(file_path):
                os.remove(file_path)
                self.status_label.config(text=f"已删除文件: {os.path.basename(file_path)}")
                # 重新扫描当前组
                self.refresh_current_group()
            else:
                messagebox.showwarning("警告", "文件不存在")
        except Exception as e:
            messagebox.showerror("错误", f"删除文件失败: {str(e)}")
            
    def delete_smaller_files(self):
        """删除当前组中除了最大文件外的所有文件"""
        if not self.similar_groups:
            messagebox.showwarning("警告", "没有相似图片组")
            return
            
        group = self.similar_groups[self.current_group_index]
        if len(group) <= 1:
            messagebox.showinfo("信息", "当前组只有一张图片，无需删除")
            return
            
        # 确认删除
        files_to_delete = group[1:]  # 除了第一个（最大的）文件
        total_size = sum(size for _, size in files_to_delete)
        
        result = messagebox.askyesno(
            "确认删除", 
            f"将删除 {len(files_to_delete)} 个文件，总大小: {self.format_file_size(total_size)}\n"
            f"保留最大文件: {os.path.basename(group[0][0])}\n\n"
            f"确定要删除吗？"
        )
        
        if result:
            deleted_count = 0
            failed_files = []
            
            for file_path, _ in files_to_delete:
                try:
                    if os.path.exists(file_path):
                        os.remove(file_path)
                        deleted_count += 1
                    else:
                        failed_files.append(file_path)
                except Exception as e:
                    failed_files.append(f"{file_path} (错误: {str(e)})")
            
            # 更新状态
            if failed_files:
                messagebox.showwarning(
                    "部分删除失败", 
                    f"成功删除 {deleted_count} 个文件\n"
                    f"失败的文件:\n" + "\n".join(failed_files)
                )
            else:
                messagebox.showinfo("删除完成", f"成功删除 {deleted_count} 个文件")
            
            # 重新扫描当前组
            self.refresh_current_group()
            
    def refresh_current_group(self):
        """刷新当前组显示"""
        if not self.similar_groups:
            return
            
        group = self.similar_groups[self.current_group_index]
        
        # 检查文件是否还存在
        existing_files = []
        for file_path, file_size in group:
            if os.path.exists(file_path):
                existing_files.append((file_path, file_size))
        
        if len(existing_files) <= 1:
            # 如果只剩一个或没有文件，从组中移除
            self.similar_groups.pop(self.current_group_index)
            
            if not self.similar_groups:
                self.status_label.config(text="所有相似图片组已处理完毕")
                self.clear_results()
                return
            else:
                # 调整当前组索引
                if self.current_group_index >= len(self.similar_groups):
                    self.current_group_index = len(self.similar_groups) - 1
        else:
            # 更新当前组
            existing_files.sort(key=lambda x: x[1], reverse=True)
            self.similar_groups[self.current_group_index] = existing_files
        
        # 重新显示当前组
        self.display_current_group()
        self.status_label.config(text=f"找到 {len(self.similar_groups)} 组相似图片")
            
    def format_file_size(self, size_bytes):
        """格式化文件大小"""
        if size_bytes == 0:
            return "0 B"
        
        size_names = ["B", "KB", "MB", "GB"]
        i = 0
        while size_bytes >= 1024 and i < len(size_names) - 1:
            size_bytes /= 1024.0
            i += 1
            
        return f"{size_bytes:.1f} {size_names[i]}"
        
    def prev_group(self):
        if self.similar_groups and self.current_group_index > 0:
            self.current_group_index -= 1
            self.display_current_group()
            
    def next_group(self):
        if self.similar_groups and self.current_group_index < len(self.similar_groups) - 1:
            self.current_group_index += 1
            self.display_current_group()
            
    def clear_results(self):
        for widget in self.scrollable_frame.winfo_children():
            widget.destroy()
            
    def show_message(self, message):
        self.status_label.config(text=message)
        messagebox.showinfo("信息", message)
        
    def show_error(self, message):
        self.status_label.config(text=message)
        messagebox.showerror("错误", message)

def main():
    root = tk.Tk()
    app = DuplicateImageFinder(root)
    root.mainloop()

if __name__ == "__main__":
    main()
